`include "mycpu.h"
module if_stage(
    input clk,
    input reset,
    //ID allowin
    input ds_allowin,
    //32: branch valid 31-0: branch target
    input [`BR_BUS_WD - 1:0] br_bus,
    //IF stage to ID stage
    output fs_to_ds_valid,
    output [`FS_TO_DS_WD - 1:0] fs_to_ds_bus,
    //exception
    input exec_flush,
    input [31:0] fs_ex_entry,
    //inst_sram
    output inst_sram_req,
    output inst_sram_wr,
    output [1:0] inst_sram_size,
    output [3:0] inst_sram_wstrb,
    output [31:0] inst_sram_addr,
    output [31:0] inst_sram_wdata,
    input inst_sram_addr_ok,
    input inst_sram_data_ok,
    input [31:0] inst_sram_rdata
);
    wire br_stall;
    wire pfs_to_fs_valid;
    wire pfs_ready_go;
    wire pfs_allowin;
    reg pfs_req_en;

    wire fs_inst_cancel;
    reg exc_flush_r_valid;
    reg [31:0] exc_entry_r;

    reg pfs_addr_ok_r;
    wire pfs_addr_ok;
    reg fs_data_ok_r;
    wire fs_data_ok;

    reg [31:0] fs_inst_buf;
    reg fs_inst_buf_valid;

    reg br_taken_r;
    reg [31:0] br_target_r;


    //fs valid bit
    wire fs_allowin;
    reg fs_throw;
    reg fs_valid;
    wire fs_ready_go;
    //pc
    reg [31:0] fs_pc;
    wire [31:0] nextpc;
    wire [31:0] seq_pc;
    wire [31:0] fs_inst;
    //exp13: ADEF
    wire fs_ex_adef;
    wire pfs_ex_adef;
    reg fs_ex_adef_r;

    assign fs_ex_adef = fs_ex_adef_r;
    assign pfs_ex_adef = nextpc[0] | nextpc[1];
    always @(posedge clk) begin
        if(reset)
            fs_ex_adef_r <= 1'b0;
        else if(inst_sram_req && pfs_addr_ok && fs_allowin)
            fs_ex_adef_r <= pfs_ex_adef;
    end

    assign fs_to_ds_bus[`FS_TO_DS_WD - 1:0] = {
        fs_ex_adef,
        fs_inst[31:0],
        fs_pc[31:0]
    };
    assign inst_sram_wr = 1'b0;
    assign inst_sram_wstrb[3:0] = {4{1'b0}};
    assign inst_sram_wdata[31:0] = 32'b0;
//pre-IF
    assign inst_sram_size[1:0] = {1'b1, 1'b0};
    assign inst_sram_req = ~reset & fs_allowin & pfs_req_en & ~br_stall;
    assign inst_sram_addr = nextpc;


    assign br_stall = br_bus[33];

    assign pfs_allowin = pfs_ready_go & fs_allowin;
    assign pfs_ready_go = pfs_addr_ok & inst_sram_req;
    assign pfs_to_fs_valid  = ~reset & pfs_ready_go;
    assign pfs_addr_ok = inst_sram_addr_ok;

    assign fs_inst_cancel = exec_flush | exc_flush_r_valid;  
    assign nextpc[31:0] = {32{exec_flush}} & fs_ex_entry[31:0]
                        | {32{exc_flush_r_valid}} & exc_entry_r[31:0]
                        | {32{~br_taken_r & br_bus[32] & ~fs_inst_cancel}} & br_bus[31:0]
                        | {32{br_taken_r & ~fs_inst_cancel}} & br_target_r[31:0]
                        | {32{~fs_inst_cancel & ~(~br_taken_r & br_bus[32])
                               & ~br_taken_r}} & seq_pc;
    assign seq_pc = fs_pc + 32'h4;

    //pfs_req_en : inst_sram_req enable bit
        always @(posedge clk) begin
            if(reset)
                pfs_req_en <= 1'b1;
            else if(inst_sram_req & pfs_addr_ok)
                pfs_req_en <= 1'b0;
            else if(fs_to_ds_valid & ds_allowin)
                pfs_req_en <= 1'b1;
            else if(fs_inst_cancel & pfs_addr_ok_r)
                pfs_req_en <= 1'b1;
        end
    //pfs_addr_ok_r: already shaked hands, wait for data_ok
        always @(posedge clk) begin
            if(reset)
                pfs_addr_ok_r <= 1'b0;
            else if(inst_sram_req & pfs_addr_ok)
                pfs_addr_ok_r <= 1'b1;
            else if(fs_data_ok)
                pfs_addr_ok_r <= 1'b0;
        end
    //br r: buffer for branch taken and branch target
        always @(posedge clk) begin
            if(reset)
                br_taken_r <= 1'b0;
            else if(fs_allowin && pfs_addr_ok)
                br_taken_r <= 1'b0;
            else if(~br_stall && br_bus[32])
                br_taken_r <= 1'b1;
        end
        always @(posedge clk) begin
            if(~br_stall && br_bus[32])
                br_target_r <= br_bus[31:0];
        end
    //ex r: buffer for exc_flush and exc_entry
        always @(posedge clk) begin
            if(reset)
                exc_flush_r_valid <= 1'b0;
            else if(pfs_to_fs_valid && fs_allowin)
                exc_flush_r_valid <= 1'b0;
            else if(exec_flush & ~(inst_sram_req & pfs_addr_ok))
                exc_flush_r_valid <= 1'b1;
        end
        always @(posedge clk) begin
            if(exec_flush)
                exc_entry_r <= fs_ex_entry;
        end
//IF stage
    assign fs_ready_go = (fs_data_ok | fs_data_ok_r) & ~fs_throw & ~fs_inst_cancel;
    assign fs_allowin = !fs_valid || fs_ready_go && ds_allowin;
    assign fs_to_ds_valid = fs_valid && fs_ready_go;
    assign fs_data_ok = inst_sram_data_ok;
    assign fs_inst = fs_inst_buf_valid ? fs_inst_buf : inst_sram_rdata;
    //fs_valid & fs_pc
        always @(posedge clk) begin
            if(reset) begin
                fs_pc <= 32'h1bfffffc;
            end
            else if(fs_allowin && pfs_to_fs_valid)begin
                fs_pc <= nextpc;
            end
        end
        always @(posedge clk) begin
            if(reset) begin
                fs_valid <= 1'b0;
            end
            else if(exec_flush && !fs_allowin) begin
                fs_valid <= 1'b0;
            end
            else if(fs_allowin)begin
                fs_valid <= pfs_to_fs_valid;
            end
        end
    //fs_data_ok_r: buffer for inst_sram_data_ok
        always @(posedge clk) begin
            if(reset)
                fs_data_ok_r <= 1'b0;
            else if(ds_allowin)
                fs_data_ok_r <= 1'b0;
            else if(fs_data_ok && !ds_allowin)
                fs_data_ok_r <= 1'b1;
        end
    //fs_inst_buf: buffer for inst from sram when ds_allowin is 1'b0
        always @(posedge clk) begin
            if(reset)
                fs_inst_buf_valid <= 1'b0;
            else if(ds_allowin || fs_inst_cancel)
                fs_inst_buf_valid <= 1'b0;
            else if(fs_data_ok && !ds_allowin && !fs_inst_cancel && !fs_throw)
                fs_inst_buf_valid <= 1'b1;
        end
        always @(posedge clk) begin
            if(fs_data_ok)
                fs_inst_buf <= inst_sram_rdata;
        end
    //fs_throw: 1'b1 when next read data in IF is invalid.
        always @(posedge clk) begin
            if(reset)
                fs_throw <= 1'b0;
            else if(fs_data_ok)
                fs_throw <= 1'b0;
            else if(fs_inst_cancel && pfs_addr_ok_r && !fs_ready_go)
                fs_throw <= 1'b1;
        end
endmodule
